package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"syscall"
	"time"

	mqtt "github.com/eclipse/paho.mqtt.golang"
	"github.com/gogo/protobuf/proto"
	"github.com/golang/snappy"
	"github.com/prometheus/common/config"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/prompb"
)

var (
	// 命令行参数
	remotePath   = flag.String("remote.path", "http://192.168.63.30:8428/api/v1/write", "Remote write Path.")
	mqttBroker   = flag.String("mqtt.broker", "mqtt://192.168.63.30:1883", "default mqtt broker")
	mqttTopic    = flag.String("mqtt.topic", "sensor/#", "default mqtt topic")
	mqttClientID = flag.String("mqtt.clientid", "go-mqtt-subscriber", "default mqtt clientid")
)

type SensorData struct {
	Info map[string]interface{} `json:"info"`
	Data Sensors                `json:"sensor"`
}

type Sensors map[string]SensorValue

type SensorValue interface{}

type Sensor struct {
	Name string `json:"name"`
	SensorValue
}

var defaultRW = NewRemoteWriter(*remotePath)

// 定义消息处理函数
var messageHandler mqtt.MessageHandler = func(client mqtt.Client, msg mqtt.Message) {
	fmt.Printf("Received message: %s from topic: %s\n", msg.Payload(), msg.Topic())
	var sd SensorData
	err := json.Unmarshal(msg.Payload(), &sd)
	if err != nil {
		log.Println(fmt.Errorf("error unmarshaling JSON: %v", err))
		return
	}
	sendSamples, _ := sd.ConvertInfoToPromPB()
	fmt.Println("pd is :", sendSamples)
	pBuf := proto.NewBuffer(nil)
	req, err := buildWriteRequest(sendSamples, nil, pBuf, nil)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println("req is : ", req)

	i, err := defaultRW.Store(context.Background(), req)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println("http code : ", i)
}

func main() {
	flag.Parse()
	// MQTT服务器连接配置
	broker := *mqttBroker
	clientID := *mqttClientID
	topic := *mqttTopic
	defaultRW = NewRemoteWriter(*remotePath)
	// 创建MQTT客户端选项
	opts := mqtt.NewClientOptions()
	opts.AddBroker(broker)
	opts.SetClientID(clientID)

	// 设置连接和消息处理回调函数
	var connectHandler mqtt.OnConnectHandler = func(client mqtt.Client) {
		fmt.Println("Connected to MQTT broker")
		if token := client.Subscribe(topic, 0, messageHandler); token.Wait() && token.Error() != nil {
			fmt.Println("Error subscribing to topic:", token.Error())
		}
	}
	opts.SetOnConnectHandler(connectHandler)

	// 创建MQTT客户端实例
	client := mqtt.NewClient(opts)

	// 连接到MQTT服务器
	if token := client.Connect(); token.Wait() && token.Error() != nil {
		log.Fatal("Error connecting to MQTT broker:", token.Error())
	}

	// 等待中断信号（如Ctrl+C）以优雅地关闭连接
	var exitSignal = make(chan os.Signal, 1)
	signal.Notify(exitSignal, syscall.SIGINT, syscall.SIGTERM)
	<-exitSignal

	// 取消订阅主题
	if token := client.Unsubscribe(topic); token.Wait() && token.Error() != nil {
		fmt.Println("Error unsubscribing from topic:", token.Error())
	}
	// 断开与MQTT服务器的连接
	client.Disconnect(250)
}

// 将interface{}类型的info数据转换为prompb格式的指标
func (s SensorData) ConvertInfoToPromPB() ([]prompb.TimeSeries, error) {
	var infoLabels []prompb.Label
	var infoMetric []prompb.TimeSeries

	info := s.Info
	sensors := s.Data
	ts := int64(0)
	utc, ok := info["utc"]
	if ok {
		switch utc := utc.(type) {
		case int64:
			ts = utc * 1000
		case float64:
			ts = int64(utc) * 1000
		}

	}
	unixNanoTimestamp := time.Now().UnixMilli()
	fmt.Println("ts is : ", ts, ", unixNanoTimestamp is : ", unixNanoTimestamp)
	if ts > unixNanoTimestamp {
		ts = unixNanoTimestamp
	}
	fmt.Println("ts new is : ", ts, ", unixNanoTimestamp is : ", unixNanoTimestamp)
	for key, value := range info {
		if key == "utc" {
			continue
		}
		infoLabels = append(infoLabels, prompb.Label{
			Name:  key,
			Value: fmt.Sprintf("%v", value),
		})
	}

	for k, v := range sensors {
		tmpLabels := []prompb.Label{}
		tmpLabels = append(tmpLabels, infoLabels...)
		switch v := v.(type) {
		case map[string]interface{}:
			for key, value := range v {
				if key != "name" {
					continue
				}
				switch value := value.(type) {
				case string:
					tmpLabels = appendLabels(infoLabels, prompb.Label{
						Name:  "sensor",
						Value: value,
					})
				}
			}
			for key, value := range v {
				tmpLabels = appendLabels(tmpLabels, prompb.Label{
					Name:  "__name__",
					Value: fmt.Sprintf("sensor_%s_%s", k, key),
				})
				var sample prompb.Sample
				switch value := value.(type) {
				case float64:
					sample = prompb.Sample{
						Value:     value,
						Timestamp: ts,
					}
				case int64:
					sample = prompb.Sample{
						Value:     float64(value),
						Timestamp: ts,
					}
				default:
					log.Println("Unsupported sensor data type for sample:", value, ts)
					continue
				}
				infoMetric = append(infoMetric, prompb.TimeSeries{
					Labels:  tmpLabels,
					Samples: []prompb.Sample{sample},
				})
			}

		default:
			log.Println("Unsupported sensor data type:", v)
			return nil, fmt.Errorf("unsupported sensor data type: %v", v)
		}

	}
	log.Println("infoMetric : ", infoMetric)
	return infoMetric, nil
}

func appendLabels(labels []prompb.Label, labels2 ...prompb.Label) []prompb.Label {
	if len(labels) == 0 {
		return labels2
	}
	var maplabels = map[string]prompb.Label{}
	for _, lal := range labels {
		maplabels[lal.Name] = lal
	}

	for _, lal := range labels2 {
		maplabels[lal.Name] = lal
	}
	var newLabels []prompb.Label
	for _, lal := range maplabels {
		newLabels = append(newLabels, lal)
	}
	sort.Slice(newLabels, func(i, j int) bool {
		return newLabels[i].Name < newLabels[j].Name
	})
	return newLabels
}

type RemoteWriter struct {
	State   uint16
	UseStat bool

	Addr    string
	Client  *http.Client
	timeout time.Duration
}

const defaultBackoff = 0
const maxErrMsgLen = 1024

type RecoverableError struct {
	error
	retryAfter model.Duration
}

func NewRemoteWriter(addr string) *RemoteWriter {
	httpClient, err := config.NewClientFromConfig(config.DefaultHTTPClientConfig, addr)
	if err != nil {
		log.Printf("http client err: %v", err)
		return nil
	}

	rt, err := config.NewRoundTripperFromConfig(config.DefaultHTTPClientConfig, addr)
	if err != nil {
		log.Printf("http client err: %v", err)
		return nil
	}
	httpClient.Transport = rt
	return &RemoteWriter{
		Addr:    addr,
		Client:  httpClient,
		timeout: 5 * time.Second,
		UseStat: true,
	}
}

func (c *RemoteWriter) Store(ctx context.Context, req []byte) (int, error) {
	httpReq, err := http.NewRequest("POST", c.Addr, bytes.NewReader(req))
	if err != nil {
		return 404, err
	}

	httpReq.Header.Add("Content-Encoding", "snappy")
	httpReq.Header.Set("Content-Type", "application/x-protobuf")
	httpReq.Header.Set("User-Agent", "vm-route")
	httpReq.Header.Set("X-Prometheus-Remote-Write-Version", "0.1.0")
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()
	httpResp, err := c.Client.Do(httpReq.WithContext(ctx))
	if err != nil {
		if nil == httpResp {
			return 500, RecoverableError{err, defaultBackoff}
		}
		return httpResp.StatusCode, RecoverableError{err, defaultBackoff}
	}

	defer func() {
		io.Copy(io.Discard, httpResp.Body)
		httpResp.Body.Close()
	}()

	if httpResp.StatusCode/100 != 2 {
		scanner := bufio.NewScanner(io.LimitReader(httpResp.Body, maxErrMsgLen))
		line := ""
		if scanner.Scan() {
			line = scanner.Text()
		}
		err = fmt.Errorf("server returned HTTP status %s: %s", httpResp.Status, line)
	}
	if httpResp.StatusCode/100 == 5 {
		return httpResp.StatusCode, RecoverableError{err, defaultBackoff}
	}

	return httpResp.StatusCode, err
}

func buildWriteRequest(samples []prompb.TimeSeries, metadata []prompb.MetricMetadata, pBuf *proto.Buffer, buf []byte) ([]byte, error) {
	fmt.Println("samples : ", samples)
	req := &prompb.WriteRequest{
		Timeseries: samples,
		Metadata:   metadata,
	}

	if pBuf == nil {
		pBuf = proto.NewBuffer(nil)
	} else {
		pBuf.Reset()
	}
	err := pBuf.Marshal(req)
	if err != nil {
		return nil, err
	}

	if buf != nil {
		buf = buf[0:cap(buf)]
	}
	compressed := snappy.Encode(buf, pBuf.Bytes())
	return compressed, nil
}
